import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers.experimental import preprocessing


# 下载数据集
# from zipfile import ZipFile
# import os
# uri = "https://storage.googleapis.com/tensorflow/tf-keras-datasets/jena_climate_2009_2016.csv.zip"
# zip_path = keras.utils.get_file(
#     origin=uri, fname="jena_climate_2009_2016.csv.zip")
# zip_file = ZipFile(zip_path)
# zip_file.extractall()


csv_path = "jena_climate_2009_2016.csv"
df = pd.read_csv(csv_path)


test_split = 0.2      # 20 % of the data will be used to test the model
train_size = int(df.shape[0] * (1 - test_split))

past = 720              # timestamps 时间戳的个数，每小时获得6个
sampling_rate = 6       # 时间戳采样频率，每小时6个变成每小时取1个
future = 72             # 预测：表示输入现在，获得的是12小时后的结果
sequence_length = int(past / sampling_rate)  # sequence 长度

# 都是下标 [)
start_train_x = 0
end_train_x = train_size
start_train_y = past + future
end_train_y = train_size + past + future

start_test_x = train_size
end_test_x = df.shape[0] - (past + future)
start_test_y = train_size + past + future
end_test_y = df.shape[0]


batch_size = 256
epochs = 3



titles = [
    "Pressure",
    "Temperature",
    "Temperature in Kelvin",
    "Temperature (dew point)",
    "Relative Humidity",
    "Saturation vapor pressure",
    "Vapor pressure",
    "Vapor pressure deficit",
    "Specific humidity",
    "Water vapor concentration",
    "Airtight",
    "Wind speed",
    "Maximum wind speed",
    "Wind direction in degrees",
]

# df的列索引的名字，除了 Date Time
feature_keys = [
    "p (mbar)",
    "T (degC)",
    "Tpot (K)",
    "Tdew (degC)",
    "rh (%)",
    "VPmax (mbar)",
    "VPact (mbar)",
    "VPdef (mbar)",
    "sh (g/kg)",
    "H2OC (mmol/mol)",
    "rho (g/m**3)",
    "wv (m/s)",
    "max. wv (m/s)",
    "wd (deg)",
]

colors = [
    "blue",
    "orange",
    "green",
    "red",
    "purple",
    "brown",
    "pink",
    "gray",
    "olive",
    "cyan",
]



print(
    "The selected parameters are:",
    ", ".join([titles[i] for i in [0, 1, 5, 7, 8, 10, 11]]),
)
# Pressure, Temperature, Saturation vapor pressure, Vapor pressure deficit, Specific humidity, Airtight, Wind speed




selected_features_index = [0, 1, 5, 7, 8, 10, 11]
selected_features = [feature_keys[i] for i in selected_features_index]
features = df[selected_features]
normalizer = preprocessing.Normalization()
normalizer.adapt(features.values)
features = normalizer(features.values)                 # 返回的tensor类型
features = pd.DataFrame(features.numpy())               # features 的行列索引都是从0开始

x_train = features.iloc[start_train_x:end_train_x][:].values
y_train = features.iloc[start_train_y:end_train_y][1].values


dataset_train = keras.preprocessing.timeseries_dataset_from_array(
    x_train,
    y_train,
    sequence_length=sequence_length,
    sampling_rate=sampling_rate,
    batch_size=batch_size,
)

x_test = features.iloc[start_test_x:end_test_x][:].values
y_test = features.iloc[start_test_y:end_test_y][1].values


dataset_test = keras.preprocessing.timeseries_dataset_from_array(
    x_test,
    y_test,
    sequence_length=sequence_length,
    sampling_rate=sampling_rate,
    batch_size=batch_size,
)

try:
    model = keras.models.load_model("path_to_my_model")
except:
    input_shape = (sequence_length, len(selected_features_index))
    inputs = keras.layers.Input(shape=input_shape)
    lstm_out = keras.layers.LSTM(32)(inputs)
    outputs = keras.layers.Dense(1)(lstm_out)

    model = keras.Model(inputs=inputs, outputs=outputs)
    # model.summary()
    # keras.utils.plot_model(model, "my_first_model_with_shape_info.png", show_shapes=True)

    path_checkpoint = "model_checkpoint.h5"
    es_callback = keras.callbacks.EarlyStopping(
        monitor="val_loss",
        min_delta=0,
        patience=5
    )

    modelckpt_callback = keras.callbacks.ModelCheckpoint(
        monitor="val_loss",
        filepath=path_checkpoint,
        verbose=1,
        save_weights_only=True,
        save_best_only=True,
    )

    model.compile(
        optimizer="adam",
        loss="mse"
    )
    # 直接用测试集传入validation，就不用evaluate
    history = model.fit(
        dataset_train,
        validation_data=dataset_test,
        epochs=epochs,
        callbacks=[es_callback, modelckpt_callback],
    )
model.save("path_to_my_model")



labels = ["History", "True Future", "Model Prediction"]
marker = [".-", "rx", "go"]
future = 12



# # 第一批的结果
# for x, y in dataset_test.take(1):
#     prediction = model.predict(x)
#     for i in range(50):
#         len_x = list(range(i, i + sequence_length))
#         len_y = sequence_length + future + i
#         plt.figure(figsize=(10, 10))
#         plt.plot(len_x, x[i, :, 1].numpy(), marker[0], label=labels[0])
#         plt.plot(len_y, y[i].numpy(), marker[1], label=labels[1])
#         plt.plot(len_y, prediction[i], marker[2], label=labels[2])
#         plt.legend()
#         plt.xlabel("Time-Step")
#         plt.show()

# 只看y和预测
for x, y in dataset_test.take(3):
    plt.figure()
    prediction = model.predict(x)
    for i in range(batch_size):
        len_y = sequence_length + future + i
        plt.plot(len_y, y[i].numpy(), marker[1], label=labels[1])
        plt.plot(len_y, prediction[i], marker[2], label=labels[2])
    plt.legend(labels)
    plt.xlabel("Time-Step")
    plt.show()